import sys

class McTree(object):
    """
    A suffix tree with direct access to leaves and root.
    """

    def __init__(self, s):
        """
        Constructs a suffix tree with McCreights algorithm over string s.
        """
        if not s.endswith('$'):
            s += '$'
        #self.s = s.lower()
        self.s = s#.strip('\n\r\t')
        self.leaves = [None for i in xrange(0, len(s))]
        self.head = list(self.leaves)

        #Adds root and the longest suffix string, s[1..n], to the tree.
        self.root = McEdge(self, 0,0)
        e = McLeaf(self, 0, len(s), 1, self.root)
        self.root.setSuffixLink(self.root)

        #for i in xrange(1, len(self.s)-1):
        #    self.__AddSuffix(i)

        #Add last leaf ($)
        #McLeaf(self, len(s)-1, len(s), len(s), self.root)

    def __AddSuffix(self, i):
        """
        Adds a suffix to the suffix tree.
        This should however not be done lightly, and done in order
        from the largest suffix to the smallest suffix.
        """
        #print '----------------'
        #print "Adding suffix %i: %s" %(i+1, self.s[i:])
        #print "head(%i)||tail(%i): %s||%s" %(i,i,self.getHead(i).getPath(), self.s[self.tail(i):])

        #print "Going from root, looking at tail", self.s[self.tail(i):]
        n = self.slowscan(self.root, i) #tail+1 = suffix function
        self.setHead(i+1, n)
        e = McLeaf(self, self.tail(i+1), len(self.s)+1, i+1, n)
        e.setSuffixLink(self.root)
        

    def slowscan(self, n, start):
        """
        Returns the node where the path from n differs from the
        substring s[start:stop].
        If strings differs on an edge, the edge is splitted in two.
        """
        #print "Slowscanning from %s over %s" %(str(n), self.s[start:])
        if not n.edges.has_key(self.s[start]):
            return n

        n = n.edges[self.s[start]]
        j = 0
        for i in xrange(start, len(self.s)):
            #print "Comparing %s(%i) to %s(%i);" %(self.s[i], i+1, self.s[n.start+j], n.start+j),
            if self.s[i] == self.s[n.start+j]:
                if n.start+j+1 == n.stop:
                    #print "Reached end of edge, choosing next edge by %s:" %self.s[i+1],
                    #Choose next edge .or return node.
                    if n.edges.has_key(self.s[i+1]):
                        n = n.edges[self.s[i+1]]
                        continue
                    else:
                        return n
                else:
                    j += 1
            else:
                j += n.start
                b = n.split(j, 'in slowscan')
                #Not quite sure way this, but it seems to make things work.
                b.setSuffixLink(self.root) 
                #print "Split found at %i, inserting new node" %j, b
                #print "We now have edges %s and %s" %(b, n)
                return b
                

    def fastscan(self, n, start, stop):
        """
        Should do something the same as slowscan, but in another way.
        n = startnode
        label is the edge whose label we search for from n
        label = v
        n = u
        """
        #print "Perform fastscan from %s over '%s' (%i:%i)" %(n, self.s[start:stop], start, stop)
        if start == stop:
            return (n, n.stop)

        if type(n) == McLeaf:
            return (n, n.start+(stop-start))
        
        try:
            n = n.edges[self.s[start]]
        except Exception, e:
            print n, start, stop
            raise e
            
        if (stop-start) > len(n)-1:
            return self.fastscan(n, start+len(n), stop)
        elif (stop-start) == len(n):
            return (n, n.stop)
        else:
            #b = n.split(n.start+(stop-start))
            #return b
            return (n, n.start+(stop-start))
            
     
    def __str__(self):
        """
        Returns a string and redundant representation of the suffixtree.
        """
        res = "Suffix tree over "#+self.s
        for i in xrange(0, len(self.leaves)):
            e = self.leaves[i]
            if e is None:
                continue
            #if e.getPath('').startswith('#Walrus'):
            s = e.getPath()
            res += '\n%i\t%s' %(i+1, s[:30])
        return res

    def getHead(self, i):
        """
        Err...?
        """
        if i == 1:
            return self.root
        if self.head[i-1] == None:
            raise Exception, "Head %i not set!" %i
        return self.head[i-1]

    def setHead(self, i, edge):
        """
        Points head(i) to a node.
        """
        self.head[i-1] = edge
             
    def tail(self, i):
        """
        Returns the tail of a suffix.
        """
        if self.getHead(i) == self.root:
            return i-1
        length = self.getHead(i).strdepth
        return length+i-1

    def setLeaf(self, i, edge):
        self.leaves[i-1] = edge
        

class AbstractEdge(object):
    """
    Contains both an edge and the node, that the edge points to.
    The outgoing edges are contained in a dictionary with the first
    letter in the edge as key.
    """

    def __init__(self, tree, start, stop):
        """
        Initilizes the edge with a node, such that the edge is labeled
        s[start:stop].
        """
        self.tree = tree
        self.start = start
        self.stop = stop
        self.parent = None
        self.suffixlink = None
        self.strdepth = 0

    def __len__(self):
        return self.stop - self.start

    def __str__(self):
        if self.stop == 0:
            return "#"
        return self.tree.s[self.start:self.stop]

    def getPath(self, dili='|'):
        """
        Returns the path from the root to the node.
        """
        n = self
        res = str(n)
        while True:#n.stop > 0:
            n = n.parent
            if n == None:# or n.stop == 0:
                break
            res = str(n) + dili + res
        return res

    def split(self, i, tag):
        """
        Splits the edge in two, connected by a node, such that the
        two edges are labeled s[start..i] and s[i..stop].
        Returns the new node.
        """
        #print "Split %i, %s" %(i, tag)
        e = McEdge(self.tree, self.start, i)
        self.parent.addChild(e)
        self.start = i
        e.addChild(self)
        e.suffixlink = self.suffixlink
        return e

    def getSuffixLink(self):
        """
        Ehh...
        """
        if self.suffixlink == None:
            raise Exception, 'Suffix link not set!'
        return self.suffixlink

    def setSuffixLink(self, node):
        self.suffixlink = node

class McEdge(AbstractEdge):
    """
    Implementation of the edge. Has a collection of children.
    """
    def __init__(self, tree, start, stop):
        AbstractEdge.__init__(self, tree, start, stop)
        self.edges = dict()

    def addChild(self, edge, tagline=""):
        """
        Adds an edge to the nodes collection of children and labels
        the edge with s[start..stop].
        """
        try:
            self.edges[self.tree.s[edge.start]] = edge
        except IndexError, e:
            print "Index out of range adding edge. \nlen(s) = %i, edge.start = %i" %(len(self.tree.s), edge.start)
            #print self.tree
            print edge.start, edge.stop
            print tagline
            print self.getPath()
            if not __file__.endswith('idle.pyw'):
                pass
                sys.exit(1)
        except BaseException, e:
            raise e
        edge.parent = self
        edge.strdepth = self.strdepth + (edge.stop - edge.start)

class McLeaf(AbstractEdge):
    def __init__(self, tree, start, stop, pos, parent):
        AbstractEdge.__init__(self, tree, start, stop)
        self.pos = pos
        parent.addChild(self, 'Creating new leaf '+str(pos))
        tree.setLeaf(pos, self)


#s = 'tatat'
#s = 'abaab'
#s = 'Mississippi'
#s='WalrusWalrus'
#s = 'the fisher and the walrus'
#s = 'The Walrus and The Carpenter, maybe The Walrus and The Carpenter, but The Walrus and The Carpenter.'
#tree = McTree(s)
#print tree
